﻿#region Usings

using System;
using System.Collections;
using System.Linq;
using System.Windows;
using System.Windows.Controls;
using System.Windows.Data;

#endregion

namespace TreeView
{
    [TemplatePart(Name = "PART_TreeList", Type = typeof(ListBox))]
    [TemplatePart(Name = "PART_Item", Type = typeof(TreeView))]
    public class TreeView : Control
    {
        #region Constants


        /// <summary>
        /// The control layout recursion limit
        /// </summary>
        public const int LayoutRecursionLimit = 15;

        private const string PartTreeItem = "PART_Item";
        private const string PartTreeList = "PART_TreeList";

        #endregion

        #region Readonly & Static Fields

        public static readonly DependencyProperty HeaderTemplateProperty =
            DependencyProperty.Register("HeaderTemplate", typeof(DataTemplate), typeof(TreeView), new PropertyMetadata(default(DataTemplate)));

        public static readonly DependencyProperty IsSelectedProperty =
            DependencyProperty.Register("IsSelected", typeof(bool), typeof(TreeView), new PropertyMetadata(default(bool), IsSelectedChanged));

        public static readonly DependencyProperty ItemsSourceProperty =
            DependencyProperty.Register("ItemsSource", typeof(IEnumerable), typeof(TreeView), new PropertyMetadata(default(IEnumerable)));

        public static readonly DependencyProperty SelectedItemProperty =
            DependencyProperty.Register("SelectedItem", typeof(object), typeof(TreeView), new PropertyMetadata(default(object), SelectedItemChanged));

        public static readonly DependencyProperty SubItemsSourceProperty =
            DependencyProperty.Register("SubItemsSource", typeof(IEnumerable), typeof(TreeView), new PropertyMetadata(default(IEnumerable)));

        #endregion

        #region Fields

        private bool _LayoutUpdateCompleted;
        private bool _isbubble;

        #endregion

        #region C'tors

        public TreeView()
        {
            DefaultStyleKey = typeof(TreeView);
        }

        #endregion

        #region Instance Properties

        public DataTemplate HeaderTemplate
        {
            get { return (DataTemplate)GetValue(HeaderTemplateProperty); }
            set { SetValue(HeaderTemplateProperty, value); }
        }

        public bool IsSelected
        {
            get { return (bool)GetValue(IsSelectedProperty); }
            set { SetValue(IsSelectedProperty, value); }
        }

        public IEnumerable ItemsSource
        {
            get { return (IEnumerable)GetValue(ItemsSourceProperty); }
            set { SetValue(ItemsSourceProperty, value); }
        }

        public object SelectedItem
        {
            get { return GetValue(SelectedItemProperty); }
            set { SetValue(SelectedItemProperty, value); }
        }

        public IEnumerable SubItemsSource
        {
            get { return (IEnumerable)GetValue(SubItemsSourceProperty); }
            set { SetValue(SubItemsSourceProperty, value); }
        }

        #endregion

        #region Instance Methods

        public override void OnApplyTemplate()
        {
            base.OnApplyTemplate();

            var lb = GetTemplateChild(PartTreeList) as ListBox;
            if (lb != null)
            {
                lb.LayoutUpdated += LbLayoutUpdated;
            }
        }
        /// <summary>
        /// Expand the control to the item
        /// </summary>
        /// <param name="item"></param>
        public void ExpandToItem(object item)
        {
            if (item != null)
            {
                var lb = GetTemplateChild(PartTreeList) as ListBox;

                if (lb != null)
                {
                    foreach (var listBoxItem in lb.Items)
                    {
                        if (listBoxItem == item)
                        {
                            Expand();
                            break;
                        }
                        //lb.UpdateLayout();
                        var container = lb.ItemContainerGenerator.ContainerFromItem(listBoxItem);
                        if (container != null)
                        {
                            var treeViewItem = ControlHelper.FindFirstElementInVisualTree<TreeViewItem>(container, PartTreeItem);
                            if (treeViewItem != null && treeViewItem.ExpandToItem(item))
                            {
                                Expand();
                                break;
                            }
                        }
                    }
                }
            }
        }

        private void Expand()
        {
            IsSelected = true;
        }

        public void Refresh()
        {
            var lb = GetTemplateChild(PartTreeList) as ListBox;
            if (lb != null)
            {
                var list = lb.Items.Select(item => lb.ItemContainerGenerator.ContainerFromItem(item))
                    .Where(container => container != null)
                    .Select(
                        container => ControlHelper.FindFirstElementInVisualTree<TreeViewItem>(container, PartTreeItem))
                    .Where(treeviewitem => treeviewitem != null);
                foreach (var treeviewitem in list)
                {
                    treeviewitem.SetBinding(TreeViewItem.HeaderTemplateProperty, new Binding { Source = HeaderTemplate });
                    var sipbinding = GetBindingExpression(SubItemsSourceProperty);
                    if (sipbinding != null)
                    {
                        treeviewitem.SetBinding(TreeViewItem.ItemsSourceProperty, sipbinding.ParentBinding);
                    }
                    treeviewitem.Deep = 0;
                    treeviewitem.Selected -= TreeviewitemSelected;
                    treeviewitem.Selected += TreeviewitemSelected;
                }
            }
        }

        #endregion

        #region Event Handling

        private void LbLayoutUpdated(object sender, EventArgs e)
        {
            if (!_LayoutUpdateCompleted)
            {
                Refresh();
            }
        }

        private void TreeviewitemSelected(object sender, object e)
        {
            _isbubble = true;
            SelectedItem = e;
        }

        #endregion

        #region Class Methods

        private static void IsSelectedChanged(DependencyObject d, DependencyPropertyChangedEventArgs e)
        {
            var control = d as TreeView;
            if (control != null && e.NewValue is bool && !(bool)e.NewValue)
            {
                var lb = control.GetTemplateChild(PartTreeList) as ListBox;
                if (lb != null)
                {
                    foreach (var treeviewitem in lb.Items
                                                  .Select(item => lb.ItemContainerGenerator.ContainerFromItem(item))
                                                  .Where(container => container != null)
                                                  .Select(container => ControlHelper.FindFirstElementInVisualTree<TreeViewItem>(container, PartTreeItem))
                                                  .Where(treeviewitem => treeviewitem != null))
                    {
                        treeviewitem.IsExpanded = false;
                    }
                }

                control.SelectedItem = null;
            }
        }

        private static void SelectedItemChanged(DependencyObject d, DependencyPropertyChangedEventArgs e)
        {
            var control = d as TreeView;
            if (control != null)
            {
                if (!control._isbubble)
                {
                    control.ExpandToItem(e.NewValue);
                }
                else
                {
                    control._isbubble = false;
                    control.Expand();
                }
            }
        }

        #endregion
    }
}